import random

import numpy as np
import pandas as pd
import torch
from scipy.io import arff
from sklearn.metrics import roc_curve, auc, roc_auc_score
from torch.utils.data import Dataset

from cnnN import CNN
from mixN import MixedModel
from mix1N import MixedModel2

# 检查GPU是否可用
if torch.cuda.is_available():
    # 将模型移动到GPU上
    device = 'cuda'
else:
    device = 'cpu'

modelName = 'cnn'
# modelName = 'mix'
# modelName = 'mix1'
# 处理数据
# file_path = 'dataSet/AEEEM/EQ.arff'
file_path = 'datasets/JDT.csv'
model_path = 'model/cnnM/model_0.753.ckpt'
# model_path = 'model/mixM/model_0.716.ckpt'
# model_path = 'model/mix1M/model_0.778.ckpt'

# 读取模型
def load_model(Model_path):
    Model = None
    if modelName == 'cnn':
        Model = CNN()
    elif modelName == 'mix':
        Model = MixedModel()
    elif modelName == 'mix1':
        Model = MixedModel2()
    checkpoint = torch.load(Model_path)
    # 加载状态字典，去除额外的前缀
    new_state_dict = {}
    for key, value in checkpoint.items():
        if key.startswith('cnn.') or key.startswith('gru.'):
            key = key[4:]  # 去除前缀
        new_state_dict[key] = value

    # 将新的状态字典加载到模型中
    Model.load_state_dict(new_state_dict)
    print(checkpoint.keys())
    # print(checkpoint.keys())
    # Model.load_state_dict(checkpoint['state_dict'])
    Model.to(device)
    return Model
class getDataSet(Dataset):
    def __init__(self, train=True, train_x=None, train_y=None, test_x=None, test_y=None, val=False, transform=None):
        self.train = train
        self.val = val
        self.transform = transform
        if self.train:
            self.dataset = train_x
            self.labels = train_y
        elif val:
            self.dataset = test_x
            self.labels = test_y
        else:
            self.dataset = test_x

    def __getitem__(self, index):
        if self.train:
            return torch.Tensor(self.dataset[index]).to(device), self.labels[index].to(device)
        elif self.val:
            return torch.Tensor(self.dataset[index]).to(device), self.labels[index].to(device)
        else:
            return torch.Tensor(self.dataset[index]).to(device)

    def __len__(self):
        return self.dataset.shape[0]
# 读取数据
def getData(File_path):
    # data, meta = arff.loadarff(File_path)
    # 读取csv
    data = pd.read_csv(File_path, header=None)
    data = pd.DataFrame(data)
    data = data.drop([0])
    DataSets = data.iloc[:, :-1]
    Labels = data.iloc[:, -1]

    DataSets = np.array(DataSets)
    Labels = np.array(Labels)
    # 标签的转换为0、1
    for i in range(len(Labels)):
        if Labels[i] == "b'clean'":
            Labels[i] = np.int32(0)
        else:
            Labels[i] = np.int32(1)
    num = DataSets.shape[0]
    over = num - (num // 81) * 81

    # 随机删除数据
    for i in range(over):
        index = random.randint(0, num - (i + 1))
        DataSets = np.delete(DataSets, index, axis=0)
        Labels = np.delete(Labels, index, axis=0)
    batch_size = 81
    DataSets = torch.Tensor(DataSets.astype(float)).to(device)
    Labels = torch.Tensor(Labels.astype(int)).to(device)
    train_dataset = getDataSet(train=True, train_x=DataSets, train_y=Labels)
    test_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)

    return test_loader

test_loader = getData(file_path)
model = load_model(model_path)


# 预测
model.eval()
num_epochs = 5
acc_list = []
total_step = len(test_loader)
if modelName == 'cnn':
    with torch.no_grad():
        for epoch in range(num_epochs):
            for i, (datasets, labels) in enumerate(test_loader):
                outputs = model(datasets)
                sm = torch.nn.Softmax(dim=1)
                outputs = sm(outputs)

                total = labels.size(0)
                _, maxN = torch.max(outputs.data, 0)
                _, minN = torch.min(outputs.data, 0)
                N = (maxN - minN) / 81
                # print(N)
                predicted = torch.Tensor([0 if i < 0 else 1 for i in N]).to(device)
                # print(predicted,labels)

                correct = (predicted == labels).sum().item()
                acc_list.append(correct / total)

                # 召回率
                TP = torch.sum((predicted == 1) & (labels == 1))
                FN = torch.sum((predicted == 0) & (labels == 1))
                recall = TP / (TP + FN)
                # 转为numpy
                labels = labels.cpu().numpy()
                predicted = predicted.cpu().numpy()
                # ROC曲线
                fpr, tpr, thresholds = roc_curve(labels, predicted)
                roc_auc = auc(fpr, tpr)
                # AUC值
                auc_score = roc_auc_score(labels, predicted)

                print('Epoch [{}/{}], Step [{}/{}], Accuracy: {:.2f}%, Recall: {:.2f}%, AUC: {:.2f}%'
                      .format(epoch + 1, num_epochs, i + 1, total_step,
                              (correct / total) * 100, recall * 100, auc_score * 100))

elif modelName == 'mix' or modelName == 'mix1':
    with torch.no_grad():
        for epoch in range(num_epochs):
            for i, (datasets, labels) in enumerate(test_loader):
                outputs = model(datasets)
                sm = torch.nn.Softmax(dim=0)
                outputs = sm(outputs)

                total = labels.size(0)
                # 平均值
                mean = torch.mean(outputs)

                total = labels.size(0)
                predicted = torch.Tensor([0 if i < mean else 1 for i in outputs.data]).to(device)

                correct = (predicted == labels).sum().item()
                acc_list.append(correct / total)

                # 召回率
                TP = torch.sum((predicted == 1) & (labels == 1))
                FN = torch.sum((predicted == 0) & (labels == 1))
                recall = TP / (TP + FN)
                # 转为numpy
                labels = labels.cpu().numpy()
                predicted = predicted.cpu().numpy()
                # ROC曲线
                fpr, tpr, thresholds = roc_curve(labels, predicted)
                roc_auc = auc(fpr, tpr)
                # AUC值
                auc_score = roc_auc_score(labels, predicted)

                print('Epoch [{}/{}], Step [{}/{}], Accuracy: {:.2f}%, Recall: {:.2f}%, AUC: {:.2f}%'
                      .format(epoch + 1, num_epochs, i + 1, total_step,
                              (correct / total) * 100, recall * 100, auc_score * 100))